"""
Loss
"""
import mindspore.nn as nn
from mindspore import ops


class CustomWithLossCell2D(nn.Cell):
    """
    2D模型的损失包装器
    将骨干网络与损失函数组合成一个可训练单元
    """
    def __init__(self, backbone, loss_fn):
        super(CustomWithLossCell2D, self).__init__(auto_prefix=False)
        self._backbone = backbone  # 骨干网络
        self._loss_fn = loss_fn    # 损失函数

    def construct(self, data, label, sij, rs_value):
        """
        前向传播计算损失
        
        Args:
            data: 输入特征
            label: 目标标签
            sij: 应变率张量
            rs_value: 雷诺应力值
        
        Returns:
            Tensor: 计算得到的损失值
        """
        output = self._backbone(data)  # 模型预测
        return self._loss_fn(output, label, sij, rs_value)  # 计算损失


class LossFunc2D(nn.Cell):
    """
    2D湍流模型的损失函数
    包含非负约束、数据拟合和雷诺应力约束
    """
    def __init__(self):
        super(LossFunc2D, self).__init__()
        self.criterion = nn.MSELoss()  # 均方误差损失
        self.abs = ops.Abs()           # 绝对值操作
        self.sqrt = ops.Sqrt()         # 平方根操作
        self.e_value = 2.71828         # 自然常数e

    def rs_error(self, pred, sij, rs_value):
        """
        计算雷诺应力预测误差
        基于Boussinesq假设：雷诺应力 = 2 * 涡粘性 * 应变率张量
        
        Args:
            pred: 模型预测的涡粘性
            sij: 应变率张量
            rs_value: 真实雷诺应力值
        
        Returns:
            Tensor: 雷诺应力预测误差
        """
        rs1 = 2.0 * pred * sij  # Boussinesq假设计算雷诺应力
        error = self.criterion(rs1, rs_value)  # 计算预测与真实值的MSE
        return error

    def construct(self, pred, label, sij, rs_value):
        """
        构建总损失函数
        
        Args:
            pred: 模型预测值
            label: 目标标签
            sij: 应变率张量
            rs_value: 雷诺应力值
        
        Returns:
            Tensor: 组合后的总损失
        """
        # 非负约束：确保预测的涡粘性非负
        lx0 = ((self.abs(pred) - pred) / 2.0).mean()  # 提取负值部分
        lx0 = lx0 ** 2  # 平方惩罚负值
        
        # 数据拟合损失：预测值与目标值的MSE
        lx2 = self.criterion(pred, label)
        
        # 雷诺应力约束损失（权重1e-2）
        lx3 = 1e-2 * self.rs_error(pred, sij, rs_value)
        
        # 组合所有损失项
        loss = lx0 + lx2 + lx3
        return loss


class LossToEval2D(nn.Cell):
    """
    2D模型验证时使用的损失函数
    与训练损失类似，但额外返回雷诺应力损失用于评估
    """
    def __init__(self):
        super(LossToEval2D, self).__init__()
        self.criterion = nn.MSELoss()
        self.abs = ops.Abs()
        self.relu = ops.ReLU()
        self.sqrt = ops.Sqrt()
        self.e_value = 2.71828

    def rs_error(self, pred, sij, rs_value):
        """计算雷诺应力误差（同上）"""
        rs1 = 2.0 * pred * sij
        error = self.criterion(rs1, rs_value)
        return error

    def construct(self, pred, label, sij, rs_value):
        """
        构建验证损失函数
        
        Returns:
            Tuple[Tensor]: 总损失和雷诺应力损失
        """
        # 非负约束
        lx0 = ((self.abs(pred) - pred) / 2.0).mean()
        lx0 = lx0 ** 2
        
        # 数据拟合损失
        lx2 = self.criterion(pred, label)
        
        # 雷诺应力约束损失
        lx3 = 1e-2 * self.rs_error(pred, sij, rs_value)
        
        # 返回总损失和雷诺应力损失
        loss = lx0 + lx2 + lx3
        return loss, lx3


class LossFunc3D(nn.Cell):
    """
    3D湍流模型的损失函数
    包含非负约束、加权数据拟合、雷诺应力约束和R2分数约束
    """
    def __init__(self, sca_min, sca_max):
        """
        初始化3D损失函数
        
        Args:
            sca_min: 数据最小值（用于归一化）
            sca_max: 数据最大值（用于归一化）
        """
        super(LossFunc3D, self).__init__()
        self.criterion = nn.MSELoss()
        self.abs = ops.Abs()
        self.sqrt = ops.Sqrt()
        self.sca_min = sca_min  # 数据最小值
        self.sca_max = sca_max  # 数据最大值
        self.e_value = 2.71828
        self.mean = ops.ReduceMean()  # 均值操作
        self.sum = ops.ReduceSum()    # 求和操作

    def rs_error(self, pred, sij, rs_value):
        """计算雷诺应力误差（同上）"""
        rs1 = 2.0 * pred * sij
        error = self.criterion(rs1, rs_value)
        return error

    def weighted_loss(self, dis, label, pred):
        """
        基于壁面距离的加权损失
        近壁面区域对湍流模型尤为重要，因此增加权重
        
        Args:
            dis: 到壁面的距离
            label: 目标标签
            pred: 预测值
        
        Returns:
            Tensor: 加权损失
        """
        loss = (dis * (pred - label) ** 2).mean()  # 距离越近权重越大
        return loss

    def r2_loss(self, label, pred):
        """
        计算R2分数损失
        R2分数衡量预测值与真实值的拟合优度
        
        Args:
            label: 目标标签
            pred: 预测值
        
        Returns:
            Tensor: R2损失（1-R2分数）
        """
        y_mean = self.mean(label)  # 标签均值
        ss_tot = self.sum((label - y_mean) ** 2)  # 总离差平方和
        ss_res = self.sum((label - pred) ** 2)   # 残差平方和
        r2_value = 1.0 - (ss_res / ss_tot)        # R2分数
        loss = 1.0 - r2_value                     # R2损失
        return loss

    def construct(self, pred, label, dis, sij, rs_value):
        """
        构建3D模型总损失函数
        
        Args:
            pred: 预测值
            label: 目标标签
            dis: 到壁面的距离
            sij: 应变率张量
            rs_value: 雷诺应力值
        
        Returns:
            Tensor: 组合后的总损失
        """
        # 非负约束
        lx0 = ((self.abs(pred) - pred) / 2.0).mean()
        lx0 = lx0 ** 2
        
        # 加权数据拟合损失（考虑壁面距离）
        lx2 = self.weighted_loss(dis, label, pred)
        
        # 雷诺应力约束损失（权重1e-2）
        lx3 = 1.0e-2 * self.rs_error(pred, sij, rs_value)
        
        # R2分数约束损失（权重1e-3）
        lx4 = 1.0e-3 * self.r2_loss(label, pred)
        
        # 组合所有损失项
        loss = lx0 + lx2 + lx3 + lx4
        return loss


class LossToEval3D(nn.Cell):
    """
    3D模型验证时使用的损失函数
    返回多个损失组件用于详细评估模型性能
    """
    def __init__(self, sca_min, sca_max):
        """初始化（同上）"""
        super(LossToEval3D, self).__init__()
        self.criterion = nn.MSELoss()
        self.abs = ops.Abs()
        self.sqrt = ops.Sqrt()
        self.sca_min = sca_min
        self.sca_max = sca_max
        self.e_value = 2.71828
        self.mean = ops.ReduceMean()
        self.sum = ops.ReduceSum()

    def rs_error(self, pred, sij, rs_value):
        """计算雷诺应力误差（同上）"""
        rs1 = 2.0 * pred * sij
        error = self.criterion(rs1, rs_value)
        return error

    def weighted_loss(self, dis, label, pred):
        """计算加权损失（同上）"""
        loss = (dis * (pred - label) ** 2).mean()
        return loss

    def r2_loss(self, label, pred):
        """计算R2损失（同上）"""
        y_mean = self.mean(label)
        ss_tot = self.sum((label - y_mean) ** 2)
        ss_res = self.sum((label - pred) ** 2)
        r2_value = 1.0 - (ss_res / ss_tot)
        loss = 1.0 - r2_value
        return loss

    def construct(self, pred, label, dis, sij, rs_value):
        """
        构建3D验证损失函数
        
        Returns:
            Tuple[Tensor]: 总损失、雷诺应力损失、R2损失和其他损失组件
        """
        # 非负约束
        lx0 = ((self.abs(pred) - pred) / 2.0).mean()
        lx0 = lx0 ** 2
        
        # 加权数据拟合损失
        lx2 = self.weighted_loss(dis, label, pred)
        
        # 雷诺应力约束损失
        lx3 = 1.0e-2 * self.rs_error(pred, sij, rs_value)
        
        # R2分数约束损失
        lx4 = self.r2_loss(label, pred)
        
        # 组合损失（R2损失权重1e-3）
        loss = lx0 + lx2 + lx3 + 1.0e-3 * lx4
        
        # 返回所有损失组件用于详细评估
        return loss, lx3, lx4, (lx0, lx2)


class CustomWithLossCell3D(nn.Cell):
    """
    3D模型的损失包装器
    将骨干网络与损失函数组合成一个可训练单元
    """
    def __init__(self, backbone, loss_fn):
        super(CustomWithLossCell3D, self).__init__(auto_prefix=False)
        self._backbone = backbone  # 骨干网络
        self._loss_fn = loss_fn    # 损失函数

    def construct(self, data, label, dis, sij, rs_value):
        """
        前向传播计算损失
        
        Args:
            data: 输入特征
            label: 目标标签
            dis: 到壁面的距离
            sij: 应变率张量
            rs_value: 雷诺应力值
        
        Returns:
            Tensor: 计算得到的损失值
        """
        output = self._backbone(data)  # 模型预测
        return self._loss_fn(output, label, dis, sij, rs_value)  # 计算损失